import unittest

import tensorflow as tf
from model.tf.densenet import ConvBlock, densenet121

HEIGHT, WIDTH, DEPTH = 224, 224, 3


class DenseNetTestCase(unittest.TestCase):

    def setUp(self) -> None:
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            try:
                tf.config.experimental.set_virtual_device_configuration(
                    gpus[0],
                    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512)])
            except RuntimeError as e:
                print(e)

    def test_conv_block(self):
        block = ConvBlock()
        block(tf.keras.layers.Input([HEIGHT, WIDTH, DEPTH]))
        self.assertEqual([None, 56, 56, 64], block.compute_output_shape([None, 56, 56, 64]))

    def test_densenet_121(self):
        densenet_121 = densenet121()
        densenet_121(tf.keras.layers.Input([HEIGHT, WIDTH, DEPTH]))
        print(densenet_121.summary())
        self.assertEqual((40, 200), densenet_121.compute_output_shape([40, HEIGHT, WIDTH, DEPTH]))
